from datasets import load_dataset
import json
import random

# 设置随机种子以确保可重复性
random.seed(42)

# 1. 加载 hk-o1aw-sft-16k 数据集（Hugging Face 官方地址）
dataset = load_dataset("hkair-lab/hk-o1aw-sft-16k", split="train")  # 主数据集为 train 分割

# 2. 转换为列表（便于打乱和分割）
data_list = list(dataset)

# 3. 查看数据集字段和示例（确认结构，可注释）
print("数据集字段：", dataset.features)
print("示例数据：", data_list[0])
# 4. 随机打乱数据
random.shuffle(data_list)

# 5. 按 9:1 分割训练集和验证集
split_idx = int(len(data_list) * 0.9)
train_data = data_list[:split_idx]
val_data = data_list[split_idx:]

# 6. 保存训练集（保留原生字段）
with open('train.jsonl', 'w', encoding='utf-8') as f:
    for item in train_data:
        # 直接提取原生字段，确保格式为 {"question":..., "think":..., "answer":...}
        json.dump({
            "question": item["prompt"],
            "think": item["thinking"],
            "answer": item["answer"]
        }, f, ensure_ascii=False)
        f.write('\n')

# 7. 保存验证集
with open('val.jsonl', 'w', encoding='utf-8') as f:
    for item in val_data:
        json.dump({
            "question": item["prompt"],
            "think": item["thinking"],
            "answer": item["answer"]
        }, f, ensure_ascii=False)
        f.write('\n')

# 输出处理结果
print(f"\nhk-o1aw-sft-16k 数据集分割完成：")
print(f"总数据量：{len(data_list)} 条")
print(f"训练集大小：{len(train_data)} 条")
print(f"验证集大小：{len(val_data)} 条")